from typing import Dict, Any, List, Optional
from ..embeddings.base import BaseEmbedding
from ..vectorstore.base import BaseVectorStore
from ..retriever.base import BaseRetriever
from ..llm.base import BaseLLM


class SimpleRAG:
    """
    简单的RAG（检索增强生成）系统实现
    """
    
    def __init__(
        self,
        embedding_model: BaseEmbedding,
        vectorstore: BaseVectorStore,
        retriever: BaseRetriever,
        llm: BaseLLM,
        k: int = 4
    ) -> None:
        """
        初始化RAG系统
        
        参数:
            embedding_model: 向量嵌入模型
            vectorstore: 向量存储
            retriever: 检索器
            llm: 语言模型
            k: 默认返回的相关文档数量
        """
        self.embedding_model = embedding_model
        self.vectorstore = vectorstore
        self.retriever = retriever
        self.llm = llm
        self.k = k

    def add_documents(
        self, 
        documents: List[str], 
        metadata: Optional[List[Dict[str, Any]]] = None
    ) -> None:
        """
        添加文档到RAG系统
        
        参数:
            documents: 文档内容列表
            metadata: 文档元数据列表
        """
        try:
            print("\n=== 添加文档开始 ===")
            print(f"文档数量: {len(documents)}")
            if metadata:
                print(f"元数据数量: {len(metadata)}")
            
            # 添加到向量存储
            self.vectorstore.add_texts(documents, metadata)
            
            print("=== 添加文档完成 ===")
            print(f"当前总文档数: {self.vectorstore.total_docs}")
            
        except Exception as e:
            print("\n=== 添加文档出错 ===")
            print(f"错误类型: {type(e).__name__}")
            print(f"错误信息: {str(e)}")
            raise

    def query(
        self,
        question: str,
        k: Optional[int] = None
    ) -> Dict[str, Any]:
        """
        处理查询请求
        
        参数:
            question: 查询问题
            k: 返回的相关文档数量，如果为None则使用默认值
            
        返回:
            包含答案、来源和相关文档的字典
        """
        try:
            print("\n=== RAG查询开始 ===")
            print(f"问题: {question}")
            print(f"k值: {k or self.k}")
            
            # 1. 检索相关文档
            k = k or self.k
            relevant_docs = self.retriever.get_relevant_documents(question, k=k)
            print(f"检索到 {len(relevant_docs)} 个相关文档")
            
            if not relevant_docs:
                print("未找到相关文档")
                return {
                    "answer": "抱歉，没有找到相关信息。",
                    "sources": [],
                    "relevant_docs": []
                }
            
            # 2. 构建提示
            context = "\n\n".join([
                f"文档 {i+1}:\n{doc['content']}"
                for i, doc in enumerate(relevant_docs)
            ])
            
            prompt = f"""基于以下信息回答问题。如果无法从提供的信息中找到答案，请说"抱歉，我无法从提供的信息中找到答案。"

相关信息:
{context}

问题: {question}

回答:"""
            
            print("\n=== 生成回答 ===")
            # 3. 生成回答
            answer = self.llm.generate(prompt)
            
            # 4. 准备返回结果
            sources = []
            for doc in relevant_docs:
                if isinstance(doc.get("metadata"), dict):
                    sources.append(doc["metadata"])
                else:
                    sources.append({})
            
            result = {
                "answer": answer,
                "sources": sources,
                "relevant_docs": [
                    {
                        "content": doc["content"],
                        "metadata": doc.get("metadata", {}),
                        "score": doc.get("score", 0.0)
                    }
                    for doc in relevant_docs
                ]
            }
            
            print("\n=== RAG查询完成 ===")
            print(f"找到文档数: {len(relevant_docs)}")
            print(f"回答长度: {len(answer)}")
            
            return result
            
        except Exception as e:
            print("\n=== RAG查询出错 ===")
            print(f"错误类型: {type(e).__name__}")
            print(f"错误信息: {str(e)}")
            raise
